package pgsql

import (
	"database/sql"
	"fmt"
	"strings"

	"github.com/example-user/go-pocket-etl/pkg/core"
	"github.com/example-user/go-pocket-etl/pkg/factory"
	// 匿名导入 PostgreSQL 驱动 (pq)，使其通过 init() 函数将自己注册到 database/sql 包中。
	_ "github.com/lib/pq"
)

// init 实现了组件的自注册。
func init() {
	factory.RegisterSink("pgsql", NewSink)
}

// Sink 实现了 core.Sink 接口，用于将数据批量写入到 PostgreSQL 数据库。
// 它通过使用 PostgreSQL 高效的 `COPY FROM STDIN` 命令来实现卓越的写入性能。
type Sink struct {
	db            *sql.DB           // 数据库连接池。
	table         string            // 目标表名。
	columnMapping map[string]string // 列映射关系 (Record 中的键 -> 数据库中的列名)。
}

// NewSink 是 pgsql.Sink 的构造函数，由工厂调用。
func NewSink() core.Sink {
	return &Sink{}
}

// Open 负责解析配置并建立数据库连接。
//
// YAML 配置示例:
//
//	sink:
//	  type: pgsql
//	  params:
//	    dsn: "postgres://user:password@localhost/dbname?sslmode=disable"
//	    table: "users"
//	    column_mapping:
//	      "id": "user_id"
//	      "email": "user_email"
//	      "status": "account_status"
func (s *Sink) Open(config map[string]interface{}) error {
	// 'dsn' 是必需配置。
	dsn, ok := config["dsn"].(string)
	if !ok || dsn == "" {
		return fmt.Errorf("pgsql sink: config is missing or has invalid 'dsn'")
	}

	// 'table' 是必需配置。
	table, ok := config["table"].(string)
	if !ok || table == "" {
		return fmt.Errorf("pgsql sink: config is missing or has invalid 'table'")
	}
	s.table = table

	// 'column_mapping' 是必需配置。
	mappingVal, ok := config["column_mapping"]
	if !ok {
		return fmt.Errorf("pgsql sink: config is missing required key 'column_mapping'")
	}

	// YAML 解析器会将 map 解析为 map[interface{}]interface{}，这里需要进行类型转换。
	mapping, ok := mappingVal.(map[interface{}]interface{})
	if !ok {
		return fmt.Errorf("pgsql sink: 'column_mapping' is not a valid map, actual type: %T", mappingVal)
	}

	if len(mapping) == 0 {
		return fmt.Errorf("pgsql sink: 'column_mapping' cannot be empty")
	}

	s.columnMapping = make(map[string]string)
	for k, v := range mapping {
		kStr, okK := k.(string)
		vStr, okV := v.(string)
		if !okK || !okV {
			return fmt.Errorf("pgsql sink: 'column_mapping' keys and values must all be strings")
		}
		s.columnMapping[kStr] = vStr
	}

	var err error
	s.db, err = sql.Open("postgres", dsn)
	if err != nil {
		return fmt.Errorf("pgsql sink: failed to initialize database connection pool: %w", err)
	}

	if err = s.db.Ping(); err != nil {
		return fmt.Errorf("pgsql sink: failed to connect to database: %w", err)
	}

	return nil
}

// Write 使用 PostgreSQL 高效的 `COPY FROM STDIN` 命令来批量写入一批记录。
//
// 设计说明:
//  1. **卓越性能**: `COPY` 是 PostgreSQL 官方推荐的、用于大批量数据加载的最高效方法。
//     它绕过了许多常规 `INSERT` 的开销，如 SQL 解析和查询计划，从而实现极高的吞吐量。
//  2. **事务性**: 整个 `COPY` 操作被包裹在一个数据库事务中，保证了批次写入的原子性。
//  3. **流式协议**: `pq` 驱动通过 `tx.Prepare` 启动 `COPY` 模式，后续的 `stmt.Exec` 调用
//     并不会执行单独的 SQL，而是将数据流式地发送到服务器，直到 `stmt.Close()` 被调用。
//  4. **关键的 `stmt.Close()`**: 调用 `stmt.Close()` 是结束 `COPY` 操作并使其生效的关键步骤。
//     如果此调用失败，整个批次的数据都可能丢失，因此必须检查其返回的错误。
func (s *Sink) Write(records []core.Record) error {
	if len(records) == 0 {
		return nil
	}

	if s.db == nil {
		return fmt.Errorf("pgsql sink: database connection is not open")
	}

	tx, err := s.db.Begin()
	if err != nil {
		return fmt.Errorf("pgsql sink: failed to begin transaction: %w", err)
	}
	defer tx.Rollback() // 确保在出错时回滚

	// 根据 column_mapping 准备数据库列名顺序。
	dbColumns := make([]string, 0, len(s.columnMapping))
	recordKeysInOrder := make([]string, 0, len(s.columnMapping))

	for recordKey, dbCol := range s.columnMapping {
		// 为列名加上双引号以处理保留字和大小写。
		dbColumns = append(dbColumns, `"`+dbCol+`"`)
		recordKeysInOrder = append(recordKeysInOrder, recordKey)
	}

	// 准备 COPY 语句。这会启动 COPY 模式。
	stmt, err := tx.Prepare(fmt.Sprintf("COPY %s (%s) FROM STDIN", s.table, strings.Join(dbColumns, ", ")))
	if err != nil {
		return fmt.Errorf("pgsql sink: failed to prepare COPY statement: %w", err)
	}
	defer stmt.Close()

	// 逐条记录将数据流式传输到服务器。
	for _, record := range records {
		args := make([]interface{}, len(recordKeysInOrder))
		for i, recordKey := range recordKeysInOrder {
			// 如果记录中缺少该键，则插入 NULL。
			args[i] = record[recordKey]
		}
		_, err = stmt.Exec(args...)
		if err != nil {
			return fmt.Errorf("pgsql sink: failed to stream data for record %v: %w", record, err)
		}
	}

	// 这是至关重要的一步：关闭 statement 会结束 COPY 操作并让服务器开始处理数据。
	// 必须检查这里的错误。
	if err := stmt.Close(); err != nil {
		return fmt.Errorf("pgsql sink: failed to finalize COPY operation: %w", err)
	}

	// 如果一切顺利，提交事务。
	return tx.Commit()
}

// Close 负责关闭数据库连接池，释放所有底层连接。
func (s *Sink) Close() error {
	if s.db != nil {
		return s.db.Close()
	}
	return nil
}
